import os
import xml.etree.ElementTree as ET
from PIL import Image

# 类别名到 YOLO class_id 的映射
CLASS_NAME_TO_ID = {
    'car': 0,
    'truck': 1,
    'pedestrian': 2,
    'traffic_cone': 3,
    'forklift': 4,
    'tractor': 5,
    'flatbed': 6,
    'trailer': 7,
    'cyclist': 8,
    'bicycle': 9,
    'water_horse': 10,
    'guardrail': 11,
    'unknown': 12
}

# 要处理的 scene-i 范围
scene_range = range(1, 23)  # 可修改范围

for i in scene_range:
    print(f"\n===> 处理 scene-{i}_image")

    SCENE_DIR = f'datasets/scene_test/scene-{i}_image'
    IMAGE_DIR = os.path.join(SCENE_DIR, 'images')
    ANNOTATION_FILE = os.path.join(SCENE_DIR, 'annotations.xml')
    LABELS_DIR = os.path.join(SCENE_DIR, 'labels')
    os.makedirs(LABELS_DIR, exist_ok=True)

    if not os.path.exists(ANNOTATION_FILE):
        print(f"[跳过] 未找到: {ANNOTATION_FILE}")
        continue

    # 解析该 scene 的 annotation 文件
    tree = ET.parse(ANNOTATION_FILE)
    root = tree.getroot()

    for image_tag in root.findall('image'):
        file_name =f"scene-{i}_"+ os.path.basename(image_tag.get('name'))
        image_path = os.path.join(IMAGE_DIR, file_name)

        if not os.path.exists(image_path):
            print(f"[跳过] 图像不存在: {file_name}")
            continue

        # 读取图像尺寸
        try:
            with Image.open(image_path) as img:
                img_width, img_height = img.size
        except:
            print(f"[错误] 无法读取图像: {file_name}")
            continue

        # 输出标签路径
        label_file = os.path.join(LABELS_DIR, os.path.splitext(file_name)[0] + '.txt')
        yolo_lines = []

        # 遍历该图像的所有标注框
        for box in image_tag.findall('box'):
            label = box.get('label')
            if label not in CLASS_NAME_TO_ID:
                print(f"[跳过] 未知类别: {label}")
                continue

            class_id = CLASS_NAME_TO_ID[label]
            xtl = float(box.get('xtl'))
            ytl = float(box.get('ytl'))
            xbr = float(box.get('xbr'))
            ybr = float(box.get('ybr'))

            x_center = (xtl + xbr) / 2 / img_width
            y_center = (ytl + ybr) / 2 / img_height
            width = (xbr - xtl) / img_width
            height = (ybr - ytl) / img_height

            yolo_line = f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
            yolo_lines.append(yolo_line)

        # 写入 YOLO 标签文件
        with open(label_file, 'w') as f:
            f.write('\n'.join(yolo_lines))

        print(f"[完成] {file_name} => labels/{os.path.basename(label_file)}")
